/***********************************************************************
 * Copyright Henry Minsky (hqm@alum.mit.edu) 1991-2009
 * (Ported to C# by Jake Tober 2010)
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 * 
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 * 
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE. 
 *
 * Commercial licensing is available under a separate license, please
 * contact author for details.
 *
 * Source code is available at http://rscode.sourceforge.net
 * Berlekamp-Peterson and Berlekamp-Massey Algorithms for error-location
 *
 * From Cain, Clark, "Error-Correction Coding For Digital Communications", pp. 205.
 *
 * This finds the coefficients of the error locator polynomial.
 *
 * The roots are then found by looking for the values of a^n
 * where evaluating the polynomial yields zero.
 *
 * Error correction is done using the error-evaluator equation  on pp 207.
 *
 */
using Logger;
using System;
using System.Diagnostics.Contracts;

namespace RSCode
{
    internal class Berlekamp
    {
        public Berlekamp(int MAXDEG, Galois galois)
        {
            Contract.Requires(galois != null);
            Contract.Requires(MAXDEG > 0);            

            Lambda = new byte[MAXDEG];
            Omega = new byte[MAXDEG];
            ErrorLocs = new int[256];
            ErasureLocs = new int[256];
            this.galois = galois;
        }

        private const int ErrorLocs_length = 256;
        private const int ErasureLocs_length = 256;

        public void Reset()
        {
            RS.zeroArray(Lambda);
            RS.zeroArray(Omega);
            RS.zeroArray(ErrorLocs);
            RS.zeroArray(ErasureLocs);
        }

        [ContractInvariantMethod]
        private void ObjectInvariant()
        {
            Contract.Invariant(galois != null);
            Contract.Invariant(Lambda != null);
            Contract.Invariant(Omega != null);
            Contract.Invariant(ErrorLocs != null);
            Contract.Invariant(ErasureLocs != null);
            Contract.Invariant(ErrorLocs.Length == ErrorLocs_length);
            Contract.Invariant(ErasureLocs.Length == ErasureLocs_length);
        }

        private Log log = Log.Create<Berlekamp>();

        /* The Error Locator Polynomial, also known as Lambda or Sigma. Lambda[0] == 1 */
        readonly byte[] Lambda;

        /* The Error Evaluator Polynomial */
       readonly byte[] Omega;

        /* error locations found using Chien's search*/
        int[] ErrorLocs;
        int NErrors;

        /* erasure flags */
        int[] ErasureLocs;
        int NErasures;

        Galois galois;

        /* From  Cain, Clark, "Error-Correction Coding For Digital Communications", pp. 216. */
        void Modified_Berlekamp_Massey(int MAXDEG, int NPAR, byte[] synBytes)
        {
            Contract.Requires(synBytes != null);
            Contract.Requires(MAXDEG > 0);           

            int n, L, L2, k, d, i;
            byte[] psi = new byte[MAXDEG];
            byte[] psi2 = new byte[MAXDEG];
            byte[] D = new byte[MAXDEG];
            byte[] gamma = new byte[MAXDEG];

            /* initialize Gamma, the erasure locator polynomial */
            init_gamma(gamma, MAXDEG);

            /* initialize to z */
            copy_poly(D, gamma, MAXDEG);
            mul_z_poly(D, MAXDEG);

            copy_poly(psi, gamma,MAXDEG);
            k = -1; L = NErasures;

            for (n = NErasures; n < NPAR; n++)
            {
                d = compute_discrepancy(psi, synBytes, L, n);
                if (d != 0)
                {
                    /* psi2 = psi - d*D */
                    for (i = 0; i < MAXDEG; i++) psi2[i] = (byte) (psi[i] ^ galois.gmult(d, D[i]));

                    if (L < (n - k))
                    {
                        L2 = n - k;
                        k = n - L;
                        /* D = scale_poly(ginv(d), psi); */
                        for (i = 0; i < MAXDEG; i++) D[i] = galois.gmult(psi[i], galois.ginv(d));
                        L = L2;
                    }
                    /* psi = psi2 */
                    for (i = 0; i < MAXDEG; i++) psi[i] = psi2[i];
                }
                mul_z_poly(D, MAXDEG);
            }
            for (i = 0; i < MAXDEG; i++) Lambda[i] = psi[i];
            compute_modified_omega(MAXDEG, NPAR, synBytes);
        }

        /* given Psi (called Lambda in Modified_Berlekamp_Massey) and synBytes,
           compute the combined erasure/error evaluator polynomial as 
           Psi*S mod z^4
          */
        void compute_modified_omega(int MAXDEG, int NPAR, byte[] synBytes)
        {
            Contract.Requires(synBytes != null);
            Contract.Requires(MAXDEG > 0);
            Contract.Requires(synBytes.Length >= MAXDEG);
            Contract.Requires(Lambda.Length >= MAXDEG);
            int i;
            byte[] product = new byte[MAXDEG * 2];

            mult_polys(product, Lambda, synBytes, MAXDEG);
            zero_poly(Omega, MAXDEG);
            for (i = 0; i < NPAR; i++) Omega[i] = product[i];

        }

        /* polynomial multiplication */
        public void mult_polys(byte[] dst, byte[] p1, byte[] p2, int MAXDEG)
        {
            Contract.Requires(dst != null);
            Contract.Requires(p1 != null);
            Contract.Requires(p2 != null);
            Contract.Requires(MAXDEG > 0);
            Contract.Requires(dst.Length >= MAXDEG*2);
            Contract.Requires(p1.Length >= MAXDEG);
            Contract.Requires(p2.Length >= MAXDEG);            

            int i, j;
            byte[] tmp1 = new byte[MAXDEG * 2];

            for (i = 0; i < (MAXDEG * 2); i++) dst[i] = 0;
            for (i = 0; i < MAXDEG; i++)
            {
                for (j = MAXDEG; j < (MAXDEG * 2); j++) tmp1[j] = 0;

                /* scale tmp1 by p1[i] */
                for (j = 0; j < MAXDEG; j++) tmp1[j] = galois.gmult(p2[j], p1[i]);
                /* and mult (shift) tmp1 right by i */
                for (j = (MAXDEG * 2) - 1; j >= i; j--) tmp1[j] = tmp1[j - i];
                for (j = 0; j < i; j++) tmp1[j] = 0;

                /* add into partial product */
                for (j = 0; j < (MAXDEG * 2); j++) dst[j] ^= tmp1[j];
            }
        }

        /* gamma = product (1-z*a^Ij) for erasure locs Ij */
        void init_gamma(byte[] gamma, int MAXDEG)
        {
            Contract.Requires(gamma != null);
            Contract.Requires(MAXDEG > 0);            
            Contract.Requires(gamma.Length >= MAXDEG);            

            int e;
            byte[] tmp = new byte[MAXDEG];

            zero_poly(gamma, MAXDEG);
            zero_poly(tmp, MAXDEG);
            gamma[0] = 1;

            for (e = 0; e < NErasures; e++)
            {
                copy_poly(tmp, gamma, MAXDEG);
                scale_poly(galois.gexp_lookup(ErasureLocs[e]), tmp, MAXDEG);
                mul_z_poly(tmp, MAXDEG);
                add_polys(gamma, tmp, MAXDEG);
            }
        }

        void compute_next_omega(int d, int[] A, int[] dst, int[] src, int MAXDEG)
        {
            Contract.Requires(A != null);
            Contract.Requires(dst != null);
            Contract.Requires(src != null);
            Contract.Requires(d >= 0);
            Contract.Requires(MAXDEG > 0);
            Contract.Requires(d < Galois.glog_length);            

            int i;
            for (i = 0; i < MAXDEG; i++)
            {
                Contract.Assume(A[i] > 0);
                dst[i] = src[i] ^ galois.gmult(d, A[i]);
            }
        }

        int compute_discrepancy(byte[] lambda, byte[] S, int L, int n)
        {
            Contract.Requires(lambda != null);
            Contract.Requires(S != null);
            Contract.Requires(L <= n);
            Contract.Requires(L >= 0);
            Contract.Requires(n >= 0);
            Contract.Requires(lambda.Length >= L);
            Contract.Requires(S.Length >= n);

            int i, sum = 0;

            for (i = 0; i <= L; i++)
                sum ^= galois.gmult(lambda[i], S[n - i]);
            return (sum);
        }

        /********** polynomial arithmetic *******************/

        public void add_polys(byte[] dst, byte[] src, int MAXDEG)
        {
            Contract.Requires(dst != null);
            Contract.Requires(src != null);
            Contract.Requires(MAXDEG > 0);
            Contract.Requires(dst.Length >= MAXDEG);
            Contract.Requires(src.Length >= MAXDEG);
            int i;
            for (i = 0; i < MAXDEG; i++) dst[i] ^= src[i];
        }

        public void copy_poly(byte[] dst, byte[] src, int MAXDEG)
        {
            Contract.Requires(dst != null);
            Contract.Requires(src != null);
            Contract.Requires(MAXDEG > 0);
            Contract.Requires(dst.Length >= MAXDEG);
            Contract.Requires(src.Length >= MAXDEG);
            int i;
            for (i = 0; i < MAXDEG; i++) dst[i] = src[i];
        }

        public void scale_poly(int k, byte[] poly, int MAXDEG)
        {
            Contract.Requires(poly != null);
            Contract.Requires(MAXDEG > 0);
            Contract.Requires(k >= 0);
            Contract.Requires(k < Galois.glog_length);
            int i;
            for (i = 0; i < MAXDEG; i++) poly[i] = galois.gmult(k, poly[i]);
        }


        public void zero_poly(byte[] poly, int MAXDEG)
        {
            Contract.Requires(poly != null);
            Contract.Requires(MAXDEG > 0);
            Contract.Requires(poly.Length >= MAXDEG);

            int i;
            for (i = 0; i < MAXDEG; i++) poly[i] = 0;
        }

        /* multiply by z, i.e., shift right by 1 */
        public void mul_z_poly(byte[] src, int MAXDEG)
        {
            Contract.Requires(src != null);
            Contract.Requires(MAXDEG > 0);
            Contract.Requires(src.Length >= MAXDEG);

            int i;
            for (i = MAXDEG - 1; i > 0; i--) src[i] = src[i - 1];
            src[0] = 0;
        }


        /* Finds all the roots of an error-locator polynomial with coefficients
         * Lambda[j] by evaluating Lambda at successive values of alpha. 
         * 
         * This can be tested with the decoder's equations case.
         */
        private void Find_Roots(int NPAR)
        {
            int sum, r, k;
            NErrors = 0;

            for (r = 1; r < 256; r++)
            {
                sum = 0;
                /* evaluate lambda at r */
                for (k = 0; k < NPAR + 1; k++)
                {
                    sum ^= galois.gmult(galois.gexp_lookup((k * r) % 255), Lambda[k]);
                }
                if (sum == 0)
                {
                    ErrorLocs[NErrors] = (255 - r); NErrors++;
                    log.Debug("Root found at r = {0}, (255-r) = {1}", r, (255 - r));
                }
            }
        }

        /* Combined Erasure And Error Magnitude Computation 
         * 
         * Pass in the codeword, its size in bytes, as well as
         * an array of any known erasure locations, along the number
         * of these erasures.
         * 
         * Evaluate Omega(actually Psi)/Lambda' at the roots
         * alpha^(-i) for error locs i. 
         *
         * Returns 1 if everything ok, or 0 if an out-of-bounds error is found
         *
         */

        public int correct_errors_erasures(byte[] codeword,
                     int csize,
                     int nerasures,
                     int[] erasures, int MAXDEG, int NPAR, byte[]synBytes)
        {
            Contract.Requires(codeword != null);
            Contract.Requires(erasures != null);
            Contract.Requires(synBytes != null);
            Contract.Requires(MAXDEG >= 0);
            Contract.Requires(nerasures <= ErasureLocs_length);

            int r, i, j;
            byte err;

            /* If you want to take advantage of erasure correction, be sure to
               set NErasures and ErasureLocs[] with the locations of erasures. 
               */
            NErasures = nerasures;
            for (i = 0; i < NErasures; i++) ErasureLocs[i] = erasures[i];

            Modified_Berlekamp_Massey(MAXDEG, NPAR, synBytes);
            Find_Roots(NPAR);


            if ((NErrors <= NPAR) && NErrors > 0)
            {

                /* first check for illegal error locs */
                for (r = 0; r < NErrors; r++)
                {
                    if (ErrorLocs[r] >= csize)
                    {                        
                        log.Error("Error loc i={0} outside of codeword length {1}", i, csize);
                        throw new Exception("decoding error");
                    }
                }

                for (r = 0; r < NErrors; r++)
                {
                    int num, denom;
                    i = ErrorLocs[r];
                    /* evaluate Omega at alpha^(-i) */

                    num = 0;
                    for (j = 0; j < MAXDEG; j++)
                        num ^= galois.gmult(Omega[j], galois.gexp_lookup(((255 - i) * j) % 255));

                    /* evaluate Lambda' (derivative) at alpha^(-i) ; all odd powers disappear */
                    denom = 0;
                    for (j = 1; j < MAXDEG; j += 2)
                    {
                        denom ^= galois.gmult(Lambda[j], galois.gexp_lookup(((255 - i) * (j - 1)) % 255));
                    }

                    err = galois.gmult(num, galois.ginv(denom));
                    log.Debug("Error magnitude {0} at loc {1}", err, csize - i);

                    codeword[csize - i - 1] ^= err;
                }
                return (1);
            }
            else
            {
                if (NErrors != 0)
                {
                    log.Error("Uncorrectable codeword");
                    throw new Exception("decoding error");
                }
                else
                {
                    return (0);
                }
            }
        }
    }
}